package ws

import (
	"sync"
	"sync/atomic"
)

type ISessionManager interface {
	Add(ISession)
	Remove(ISession)

	GetSession(int64) ISession
	ForEachSession(func(ISession) bool)
	SessionCount() int
	CloseAllSession()

	SetIDSeed(seed int64)
}

type SessionManager struct {
	sesIdSeed int64
	count     int64
	sessById  sync.Map
}

func NewSessionManager(seed int64) ISessionManager {
	return &SessionManager{sesIdSeed: seed, count: 0}
}

func (s *SessionManager) SetIDSeed(seed int64) {
	atomic.StoreInt64(&s.sesIdSeed, seed)
}

func (s *SessionManager) Add(ses ISession) {
	id := atomic.AddInt64(&s.sesIdSeed, 1)

	ses.SetID(id)

	atomic.AddInt64(&s.count, 1)
	s.sessById.Store(id, ses)
}

func (s *SessionManager) Remove(ses ISession) {
	s.sessById.Delete(ses.ID())

	atomic.AddInt64(&s.count, -1)
}

func (s *SessionManager) GetSession(id int64) ISession {
	if v, ok := s.sessById.Load(id); ok {
		return v.(ISession)
	}

	return nil
}

func (s *SessionManager) ForEachSession(callback func(ISession) bool) {
	s.sessById.Range(func(key, value interface{}) bool {
		return callback(value.(ISession))
	})
}

func (s *SessionManager) CloseAllSession() {
	s.ForEachSession(func(ses ISession) bool {
		ses.Close()
		return true
	})
}

func (s *SessionManager) SessionCount() int {
	v := atomic.LoadInt64(&s.count)
	return int(v)
}
